import torch

# 1、解析数据
import numpy as np
import re

ff = open('../data/housing.data').readlines()  # 读取所有行
data = []

for item in ff:
    out = re.sub(r'\s{2,}', ' ', item).strip()  # 将多空格合并为一个空格,并去掉换行符
    # print(out)
    data.append(out.split(' '))

data = np.array(data).astype(np.float32)  # 转换为numpy类型
# print(笔记.md.shape) # (506, 14)

# 将数据切分为x,y
x = data[:, 0:-1]
y = data[:, -1]
# print(x.shape, y.shape) # (506, 13) (506,)

# 划分训练集和测试集
x_train = x[0:496, ...]
y_train = y[0:496, ...]
x_test = x[496:, ...]
y_test = y[496:, ...]


# print(x_train.shape, y_train.shape, x_test.shape, y_test.shape)
# (496, 13) (496,) (10, 13) (10,)


# 2、搭建神经网络
class Net(torch.nn.Module):

    def __init__(self, n_feature, n_output):
        super(Net, self).__init__()
        # 加入一个隐藏层
        self.hidden = torch.nn.Linear(n_feature, 100)
        # 线性回归模型
        self.predict = torch.nn.Linear(100, n_output)

    def forward(self, x):
        out = self.hidden(x)
        out = torch.relu(out)  # 加入非线性的运算
        out = self.predict(out)
        return out


net = Net(13, 1)  # 输入特征数量，输出特征数量

# 3、定义损失函数
loss_func = torch.nn.MSELoss()  # 使用均方差损失函数

# 4、定义优化器
# 随机梯度下降法,传入参数和学习率
# optimizer = torch.optim.SGD(net.parameters(), lr=0.0001)
# Adam优化器是一种梯度下降算法的变体，用于更新神经网络的权重。
# 它结合了随机梯度下降算法（SGD）和自适应学习率算法，能够快速收敛并且减少训练时间。
# Adam优化器计算出每个参数的独立自适应学习率，不需要手动调整学习率的大小，因此在实践中被广泛使用。
optimizer = torch.optim.Adam(net.parameters(), lr=0.01)

# 5、训练模型，获取参数
for i in range(10000):
    x_data = torch.tensor(x_train, dtype=torch.float32)
    y_data = torch.tensor(y_train, dtype=torch.float32)

    # 前向推导
    pred = net.forward(x_data)
    # print(pred.shape, y_data.shape) # torch.Size([496, 1]) torch.Size([496])
    # 维度不对，需要调整
    pred = torch.squeeze(pred)
    # print(pred.shape, y_data.shape)  # torch.Size([496]) torch.Size([496])
    loss = loss_func(pred, y_data) * 0.001

    # 将神经网络中的梯度置为0
    optimizer.zero_grad()
    # 进行反向传播
    loss.backward()
    # 迭代梯度
    optimizer.step()

    # 打印迭代次数和loss的变化
    print('ite:{}, loss:{}'.format(i, loss))
    # print(pred[0:10])
    # print(y_data[0:10])

    # 6、预测
    x_data = torch.tensor(x_test, dtype=torch.float32)
    y_data = torch.tensor(y_test, dtype=torch.float32)

    # 前向推导
    pred = net.forward(x_data)
    # print(pred.shape, y_data.shape) # torch.Size([496, 1]) torch.Size([496])
    # 维度不对，需要调整
    pred = torch.squeeze(pred)
    # print(pred.shape, y_data.shape)  # torch.Size([496]) torch.Size([496])
    loss_test = loss_func(pred, y_data) * 0.001
    print('ite:{}, loss_test:{}'.format(i, loss_test))

# 7、将模型进行保存
torch.save(net, '../model/model_k1.pkl') # 保存模型
# torch.load('')
torch.save(net.state_dict(), '../model/params_k1.pkl')  # 保存模型参数
# net.load_state_dict('')
